{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JndnmDMp66FL"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Dbh3a1h0zmJ4"
      },
      "source": [
        "# Train your own Keyword Spotting Model.\n",
        "[Open in Google  Colab](https://colab.research.google.com/github/google-research/google-research/blob/master/speech_embedding/speech_commands.ipynb)\n",
        "\n",
        "Before running any cells please enable GPUs for this notebook to speed it up. \n",
        "\n",
        "* *Edit* → *Notebook Settings*\n",
        "* select *GPU* from the *Hardware Accelerator* drop-down\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "7ZVCzvqWuqFV"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "%tensorflow_version 1.x\n",
        "from __future__ import division\n",
        "\n",
        "import collections\n",
        "import IPython\n",
        "import functools\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import io\n",
        "import os\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import random\n",
        "import scipy.io.wavfile\n",
        "import tarfile\n",
        "import time\n",
        "import sys\n",
        "\n",
        "from google.colab import output\n",
        "from google.colab import widgets\n",
        "from base64 import b64decode\n",
        "\n",
        "!pip install ffmpeg-python\n",
        "import ffmpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3TTWa8GN0KHo"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions and classes\n",
        "def normalized_read(filename):\n",
        "  \"\"\"Reads and normalizes a wavfile.\"\"\"\n",
        "  _, data = scipy.io.wavfile.read(open(filename, mode='rb'))\n",
        "  samples_99_percentile = np.percentile(np.abs(data), 99.9)\n",
        "  normalized_samples = data / samples_99_percentile\n",
        "  normalized_samples = np.clip(normalized_samples, -1, 1)\n",
        "  return normalized_samples\n",
        "\n",
        "class EmbeddingDataFileList(object):\n",
        "  \"\"\"Container that loads audio, stores it as embeddings and can\n",
        "  rebalance it.\"\"\"\n",
        "\n",
        "  def __init__(self, filelist,\n",
        "               data_dest_dir,\n",
        "               targets=None,\n",
        "               label_max=10000,\n",
        "               negative_label=\"negative\",\n",
        "               silence_label=\"silence\",\n",
        "               negative_multiplier=25,\n",
        "               target_samples=32000,\n",
        "               progress_bar=None,\n",
        "               embedding_model=None):\n",
        "    \"\"\"Creates an instance of `EmbeddingDataFileList`.\"\"\"\n",
        "    self._negative_label = negative_label\n",
        "    self._silence_label = silence_label\n",
        "    self._data_per_label = collections.defaultdict(list)\n",
        "    self._labelcounts = {}\n",
        "    self._label_list = targets\n",
        "    total_examples = sum([min(len(x), label_max) for x in filelist.values()])\n",
        "    total_examples -= min(len(filelist[negative_label]), label_max)\n",
        "    total_examples -= min(len(filelist[silence_label]), label_max)\n",
        "    total_examples += min(len(filelist[negative_label]), negative_multiplier * label_max)\n",
        "    total_examples += min(len(filelist[silence_label]), negative_multiplier * label_max)\n",
        "\n",
        "    print(\"loading %d examples\" % total_examples)\n",
        "    example_count = 0\n",
        "    for label in filelist:\n",
        "      if label not in self._label_list:\n",
        "        raise ValueError(\"Unknown label:\", label)\n",
        "      label_files = filelist[label]\n",
        "      random.shuffle(label_files)\n",
        "      if label == negative_label or label == silence_label:\n",
        "        multplier = negative_multiplier\n",
        "      else:\n",
        "        multplier = 1\n",
        "      for wav_file in label_files[:label_max * multplier]:\n",
        "        data = normalized_read(os.path.join(data_dest_dir, wav_file))\n",
        "        required_padding = target_samples - data.shape[0]\n",
        "        if required_padding \u003e 0:\n",
        "          data = np.pad(data, (required_padding, required_padding), 'constant')\n",
        "        self._labelcounts[label] = self._labelcounts.get(label, 0) + 1\n",
        "        if embedding_model:\n",
        "           data = embedding_model.create_embedding(data)[0][0,:,:,:]\n",
        "        self._data_per_label[label].append(data)\n",
        "        if progress_bar is not None:\n",
        "          example_count += 1\n",
        "          progress_bar.update(progress(100 * example_count/total_examples))\n",
        "\n",
        "  @property\n",
        "  def labels(self):\n",
        "    return self._label_list\n",
        "\n",
        "  def get_label(self, idx):\n",
        "    return self.labels.index(idx)\n",
        "\n",
        "  def _get_filtered_data(self, label, filter_fn):\n",
        "    idx = self.labels.index(label)\n",
        "    return [(filter_fn(x), idx) for x in self._data_per_label[label]]\n",
        "\n",
        "  def _multply_data(self, data, factor):\n",
        "    samples = int((factor - math.floor(factor)) * len(data))\n",
        "    return int(factor) * data + random.sample(data, samples)\n",
        "\n",
        "  def full_rebalance(self, negatives, labeled):\n",
        "    \"\"\"Rebalances for a given ratio of labeled to negatives.\"\"\"\n",
        "    negative_count = self._labelcounts[self._negative_label]\n",
        "    labeled_count = sum(self._labelcounts[key]\n",
        "                        for key in self._labelcounts.keys()\n",
        "                        if key not in [self._negative_label, self._silence_label])\n",
        "    labeled_multiply = labeled * negative_count / (negatives * labeled_count)\n",
        "    for label in self._data_per_label:\n",
        "      if label in [self._negative_label, self._silence_label]:\n",
        "        continue\n",
        "      self._data_per_label[label] = self._multply_data(\n",
        "          self._data_per_label[label], labeled_multiply)\n",
        "      self._labelcounts[label] = len(self._data_per_label[label])\n",
        "\n",
        "  def get_all_data_shuffled(self, filter_fn):\n",
        "    \"\"\"Returns a shuffled list containing all the data.\"\"\"\n",
        "    return self.get_all_data(filter_fn, shuffled=True)\n",
        "\n",
        "  def get_all_data(self, filter_fn, shuffled=False):\n",
        "    \"\"\"Returns a list containing all the data.\"\"\"\n",
        "    data = []\n",
        "    for label in self._data_per_label:\n",
        "      data += self._get_filtered_data(label, filter_fn)\n",
        "    if shuffled:\n",
        "      random.shuffle(data)\n",
        "    return data\n",
        "\n",
        "def cut_middle_frame(embedding, num_frames, flatten):\n",
        "  \"\"\"Extrats the middle frames for an embedding.\"\"\"\n",
        "  left_context = (embedding.shape[0] - num_frames) // 2\n",
        "  if flatten:\n",
        "    return embedding[left_context:left_context+num_frames].flatten()\n",
        "  else:\n",
        "    return embedding[left_context:left_context+num_frames]\n",
        "\n",
        "\n",
        "def progress(value, maximum=100):\n",
        "  return IPython.display.HTML(\"\"\"\n",
        "  \u003cprogress value='{value}' max='{max}' style='width: 80%'\u003e{value}\u003c/progress\u003e\n",
        "    \"\"\".format(value=value, max=maximum))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "uGjxofKb07bk"
      },
      "outputs": [],
      "source": [
        "#@title HeadTrainerClass and head model functions\n",
        "\n",
        "def _fully_connected_model_fn(embeddings, num_labels):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  net = tf.layers.flatten(embeddings)\n",
        "  logits = tf.compat.v1.layers.dense(net, num_labels, activation=None)\n",
        "  return logits\n",
        "\n",
        "framework = tf.contrib.framework\n",
        "layers = tf.contrib.layers\n",
        "\n",
        "def _conv_head_model_fn(embeddings, num_labels, context):\n",
        "  \"\"\"Builds the head model and adds a fully connected output layer.\"\"\"\n",
        "  activation_fn = tf.nn.elu\n",
        "  normalizer_fn = functools.partial(\n",
        "      layers.batch_norm, scale=True, is_training=True)\n",
        "  with framework.arg_scope([layers.conv2d], biases_initializer=None,\n",
        "                           activation_fn=None, stride=1, padding=\"SAME\"):\n",
        "    net = embeddings\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [2, 1], stride=[2, 1], padding=\"VALID\")\n",
        "    context //= 2\n",
        "    net = layers.conv2d(net, 96, [3, 1])\n",
        "    net = normalizer_fn(net)\n",
        "    net = activation_fn(net)\n",
        "    net = layers.max_pool2d(net, [context, net.shape[2]], padding=\"VALID\")\n",
        "  net = tf.layers.flatten(net)\n",
        "  logits = layers.fully_connected(\n",
        "      net, num_labels, activation_fn=None)\n",
        "  return logits\n",
        "\n",
        "class HeadTrainer(object):\n",
        "  \"\"\"A tensorflow classifier to quickly train and test on embeddings.\n",
        "\n",
        "  Only use this if you are training a very small model on a very limited amount\n",
        "  of data. If you expect the training to take any more than 15 - 20 min then use\n",
        "  something else.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, model_fn, input_shape, num_targets,\n",
        "               head_learning_rate=0.001, batch_size=64):\n",
        "    \"\"\"Creates a `HeadTrainer`.\n",
        "\n",
        "    Args:\n",
        "      model_fn: function that builds the tensorflow model, defines its loss\n",
        "          and returns the tuple (predictions, loss, accuracy).\n",
        "      input_shape: describes the shape of the models input feature.\n",
        "          Does not include a the batch dimension.\n",
        "      num_targets: Target number of keywords.\n",
        "    \"\"\"\n",
        "    self._input_shape = input_shape\n",
        "    self._output_dim = num_targets\n",
        "    self._batch_size = batch_size\n",
        "    self._graph = tf.Graph()\n",
        "    with self._graph.as_default():\n",
        "      self._feature = tf.placeholder(tf.float32, shape=([None] + input_shape))\n",
        "      self._labels = tf.placeholder(tf.int64, shape=(None))\n",
        "      module_spec = hub.create_module_spec(\n",
        "          module_fn=self._get_headmodule_fn(model_fn, num_targets))\n",
        "      self._module = hub.Module(module_spec, trainable=True)\n",
        "      logits = self._module(self._feature)\n",
        "      self._predictions = tf.nn.softmax(logits)\n",
        "      self._loss, self._accuracy = self._get_loss(\n",
        "          logits, self._labels, self._predictions)\n",
        "      self._update_weights = tf.train.AdamOptimizer(\n",
        "          learning_rate=head_learning_rate).minimize(self._loss)\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._sess.as_default():\n",
        "      with self._graph.as_default():\n",
        "        self._sess.run(tf.local_variables_initializer())\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  def _get_headmodule_fn(self, model_fn, num_targets):\n",
        "    \"\"\"Wraps the model_fn in a tf hub module.\"\"\"\n",
        "    def module_fn():\n",
        "      embeddings = tf.placeholder(\n",
        "          tf.float32, shape=([None] + self._input_shape))\n",
        "      logit = model_fn(embeddings, num_targets)\n",
        "      hub.add_signature(name='default', inputs=embeddings, outputs=logit)\n",
        "    return module_fn\n",
        "\n",
        "\n",
        "  def _get_loss(self, logits, labels, predictions):\n",
        "    \"\"\"Defines the model's loss and accuracy.\"\"\"\n",
        "    xentropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "        logits=logits, labels=labels)\n",
        "    loss = tf.reduce_mean(xentropy_loss)\n",
        "    accuracy = tf.contrib.metrics.accuracy(tf.argmax(predictions, 1), labels)\n",
        "    return loss, accuracy\n",
        "\n",
        "  def save_head_model(self, save_directory):\n",
        "    \"\"\"Saves the model.\"\"\"\n",
        "    with self._graph.as_default():\n",
        "      self._module.export(save_directory, self._sess)\n",
        "\n",
        "\n",
        "  def _feature_transform(self, batch_features, batch_labels):\n",
        "    \"\"\"Transforms lists of features and labels into into model inputs.\"\"\"\n",
        "    return np.stack(batch_features), np.stack(batch_labels)\n",
        "\n",
        "  def _batch_data(self, data, batch_size=None):\n",
        "    \"\"\"Splits the input data into batches.\"\"\"\n",
        "    batch_features = []\n",
        "    batch_labels = []\n",
        "    batch_size = batch_size or len(data)\n",
        "    for feature, label in data:\n",
        "      if feature.shape != tuple(self._input_shape):\n",
        "        raise ValueError(\n",
        "            \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "                feature.shape, self._input_shape))\n",
        "      if not 0 \u003c= label \u003c self._output_dim:\n",
        "\n",
        "        raise ValueError('Label value ({}) outside of target range'.format(\n",
        "            label))\n",
        "      batch_features.append(feature)\n",
        "      batch_labels.append(label)\n",
        "      if len(batch_features) == batch_size:\n",
        "        yield self._feature_transform(batch_features, batch_labels)\n",
        "        del batch_features[:]\n",
        "        del batch_labels[:]\n",
        "    if batch_features:\n",
        "      yield self._feature_transform(batch_features, batch_labels)\n",
        "\n",
        "  def epoch_train(self, data, epochs=1, batch_size=None):\n",
        "    \"\"\"Trains the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      epochs: Number of times this data should be trained on.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for _ in range(epochs):\n",
        "      for features, labels in self._batch_data(data, batch_size):\n",
        "        loss, accuracy, _ = self._sess.run(\n",
        "            [self._loss, self._accuracy, self._update_weights],\n",
        "            feed_dict={self._feature: features, self._labels: labels})\n",
        "        accuracy_list.append(accuracy)\n",
        "    return (sum(accuracy_list))/len(accuracy_list), loss\n",
        "\n",
        "  def test(self, data, batch_size=None):\n",
        "    \"\"\"Evaluates the model on the provided data.\n",
        "\n",
        "    Args:\n",
        "      data: List of tuples (feature, label) where feature is a np array of\n",
        "          shape `self._input_shape` and label an int less than self._output_dim.\n",
        "      batch_size: Number of feature, label pairs per batch. Overwrites\n",
        "          `self._batch_size` when set.\n",
        "\n",
        "    Returns:\n",
        "      tuple of accuracy, loss;\n",
        "          accuracy: Average training accuracy.\n",
        "          loss: Loss of the final batch.\n",
        "    \"\"\"\n",
        "    batch_size = batch_size or self._batch_size\n",
        "    accuracy_list = []\n",
        "    for features, labels in self._batch_data(data, batch_size):\n",
        "      loss, accuracy = self._sess.run(\n",
        "          [self._loss, self._accuracy],\n",
        "          feed_dict={self._feature: features, self._labels: labels})\n",
        "      accuracy_list.append(accuracy)\n",
        "    return sum(accuracy_list)/len(accuracy_list), loss\n",
        "\n",
        "  def infer(self, example_feature):\n",
        "    \"\"\"Runs inference on example_feature.\"\"\"\n",
        "    if example_feature.shape != tuple(self._input_shape):\n",
        "      raise ValueError(\n",
        "          \"Feature shape ({}) doesn't match model shape ({})\".format(\n",
        "              example_feature.shape, self._input_shape))\n",
        "    return self._sess.run(\n",
        "        self._predictions,\n",
        "        feed_dict={self._feature: np.expand_dims(example_feature, axis=0)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "QLccWY0bIJs1"
      },
      "outputs": [],
      "source": [
        "#@title TfHubWrapper Class\n",
        "\n",
        "class TfHubWrapper(object):\n",
        "  \"\"\"A loads a tf hub embedding model.\"\"\"\n",
        "  def __init__(self, embedding_model_dir):\n",
        "    \"\"\"Creates a `SavedModelWraper`.\"\"\"\n",
        "    self._graph = tf.Graph()\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._graph.as_default():\n",
        "      with self._sess.as_default():\n",
        "        module_spec = hub.load_module_spec(embedding_model_dir)\n",
        "        embedding_module = hub.Module(module_spec)\n",
        "        self._samples = tf.placeholder(\n",
        "            tf.float32, shape=[1, None], name='audio_samples')\n",
        "        self._embedding = embedding_module(self._samples)\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "    print(\"Embedding model loaded, embedding shape:\", self._embedding.shape)\n",
        "\n",
        "  def create_embedding(self, samples):\n",
        "    samples = samples.reshape((1, -1))\n",
        "    output = self._sess.run(\n",
        "        [self._embedding],\n",
        "        feed_dict={self._samples: samples})\n",
        "    return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dnttvMi9z8ed"
      },
      "source": [
        "## Load the embedding model\n",
        "\n",
        "The following info messages can be ignored\n",
        "\n",
        "\u003e *INFO:tensorflow:Saver not created because there are no variables in the graph to restore*\n",
        "\n",
        "Don't worry tf hub is restoring all the variables.\n",
        "\n",
        "You can test the model by having it produce an embedding on zeros:\n",
        "\n",
        "\n",
        "```\n",
        "speech_embedding_model.create_embedding(np.zeros((1,66000)))\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "CVBtPzmLz8ef"
      },
      "outputs": [],
      "source": [
        "embedding_model_url = \"https://tfhub.dev/google/speech_embedding/1\"\n",
        "speech_embedding_model = TfHubWrapper(embedding_model_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "R-x8ReAxH-GT"
      },
      "source": [
        "## Get and load the test data\n",
        "\n",
        "The following cell are responsible for getting the data into the colab and creating the embeddings on top which the model is trained.\n",
        "\n",
        "To train a model on a different source of data, replace the next cell with one that copies in your data and change the file scanning cell to scan it correctly.\n",
        "\n",
        "Finally, ensure that global variable MODEL_LABELS is appropriatly set.\n",
        "\n",
        "File scanning is performed to create 2 lists of wav files:\n",
        " * A training file list containing all possible training files. (All files not in testing_list.txt or validation_list.txt)\n",
        " * An evaluation file list that we will use for testing (validation_list.txt)\n",
        "\n",
        "\n",
        "File lists are actually dictionaries with the following structure:\n",
        "\n",
        "```\n",
        "{'keyword1': ['path/to/word1/example1.wav', path/to/word1/example2.wav'],\n",
        " 'keyword2': ['path/to/word2/example1.wav', path/to/word2/example2.wav'],\n",
        " ...\n",
        " 'negativ': ['path/to/negativ_example1.wav', path/to/negativ_example2.wav']}\n",
        " ```\n",
        "\n",
        "The subsequent cells assume that the file lists are stored in the  variables: *all_eval_example_files* and *all_train_example_files*.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "aZqyHmyKxFPN"
      },
      "outputs": [],
      "source": [
        "#@title Download and extract the speech commands data set\n",
        "data_source = \"http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz\"\n",
        "data_dest_dir = \"speech_commands_v0.02\"\n",
        "test_list = data_dest_dir + \"/testing_list.txt\"\n",
        "valid_list = data_dest_dir + \"/validation_list.txt\"\n",
        "\n",
        "TARGET_WORDS = 'yes,no,up,down,left,right,on,off,stop,go'\n",
        "ALL_WORDS = 'backward,bed,bird,cat,dog,down,eight,five,follow,forward,four,go,' + 'happy,house,learn,left,marvin,nine,no,off,on,one,right,seven,sheila,six,stop,' + 'three,tree,two,up,visual,wow,yes,zero'\n",
        "\n",
        "# Note: This example colab doesn't train the silence output. \n",
        "MODEL_LABELS = ['negative', 'silence'] + TARGET_WORDS.split(',')\n",
        "\n",
        "!wget http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz\n",
        "\n",
        "print(\"extracting tar archive.. this may take a few minutes.\")\n",
        "\n",
        "if not os.path.exists(data_dest_dir):\n",
        "  os.makedirs(data_dest_dir)\n",
        "#tarfile.open(\"speech_commands_v0.02.tar.gz\", 'r:gz').extractall(data_dest_dir)\n",
        "file_count = 0\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "\n",
        "with tarfile.open(\"speech_commands_v0.02.tar.gz\", 'r:gz') as speech_commands_tar:\n",
        "  for member_info in speech_commands_tar.getmembers():\n",
        "    if file_count % 100 == 0:\n",
        "      progress_bar.update(progress(100 * file_count/105800))\n",
        "    speech_commands_tar.extract(member_info, data_dest_dir)\n",
        "    file_count+=1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "C4eL0zlLYqDL"
      },
      "outputs": [],
      "source": [
        "#@title Optional: add background silence data.\n",
        "#@markdown Run this tab if you want to reduce the number of false dectetions\n",
        "#@markdown when no speech is present. \n",
        "#@markdown It shouldn't affect the eval acuracy but may increase the training\n",
        "#@markdown the data loading and training time.  \n",
        "background_dir = os.path.join(data_dest_dir, \"_background_noise_\")\n",
        "silence_dir = os.path.join(data_dest_dir, \"silence\")\n",
        "if not os.path.exists(silence_dir):\n",
        "  os.makedirs(silence_dir)\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "\n",
        "\n",
        "noises = [\"doing_the_dishes\", \"exercise_bike\", \"white_noise\", \"dude_miaowing\",\n",
        "          \"pink_noise\", \"running_tap\"]\n",
        "\n",
        "overlap = 8000\n",
        "window_size = 32000\n",
        "scales = [0.001, 0.0031, 0.01, 0.031, 0.05, 0.1, 0.31, 0.5, 0.8, 1]\n",
        "\n",
        "total_silence = len(scales) * 6 * 65 * 16000 / 8000\n",
        "silence_count = 0\n",
        "\n",
        "for scale in scales:\n",
        "  for noise in noises:\n",
        "    noise_file = os.path.join(background_dir, noise + \".wav\")\n",
        "    noise_samples = normalized_read(noise_file)\n",
        "    position = 0\n",
        "    while position + window_size \u003c= noise_samples.shape[0]:\n",
        "      windowed_samples = noise_samples[position:position+window_size] * scale\n",
        "      file_name = \"%s_%s_%s.wav\" % (noise, position, scale)\n",
        "      output_path =  os.path.join(silence_dir, file_name)\n",
        "      scipy.io.wavfile.write(output_path, 16000, windowed_samples)\n",
        "      position += overlap\n",
        "      silence_count += 1\n",
        "      progress_bar.update(progress(100 * silence_count/total_silence))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "6t9I1VXFDyfG"
      },
      "outputs": [],
      "source": [
        "#@title Scan files\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "print(\"loading filelists from: %s \" % data_dest_dir)\n",
        "\n",
        "def get_train_test_valid_split(word):\n",
        "  word_dir = os.path.join(data_dest_dir, word)\n",
        "  all_word_files = [os.path.join(word, f) for f in os.listdir(word_dir) if os.path.isfile(os.path.join(word_dir, f))]\n",
        "  word_train_files = [f for f in all_word_files if f not in test_files and f not in valid_files]\n",
        "  word_test_files = [f for f in all_word_files if f in test_files]\n",
        "  word_valid_files = [f for f in all_word_files if f in valid_files]\n",
        "  random.shuffle(word_train_files)\n",
        "  random.shuffle(word_test_files)\n",
        "  random.shuffle(word_valid_files)\n",
        "  return word_train_files, word_test_files, word_valid_files\n",
        "\n",
        "test_files = [line.rstrip() for line in open(test_list, encoding=\"ISO-8859-1\")]\n",
        "valid_files = [line.rstrip() for line in open(valid_list, encoding=\"ISO-8859-1\")]\n",
        "\n",
        "all_train_example_files = collections.defaultdict(list)\n",
        "all_eval_example_files = collections.defaultdict(list)\n",
        "\n",
        "silence_dir = os.path.join(data_dest_dir, \"silence\")\n",
        "if os.path.exists(silence_dir):\n",
        "  all_word_list = ALL_WORDS.split(',') + [\"silence\"]\n",
        "else:\n",
        "  all_word_list = ALL_WORDS.split(',')\n",
        " \n",
        "word_count = 0\n",
        "\n",
        "for word in all_word_list:\n",
        "  if word in MODEL_LABELS:\n",
        "    label = word\n",
        "  else:\n",
        "    label = \"negative\"\n",
        "  train_files, eval_files, _ = get_train_test_valid_split(word)\n",
        "  all_train_example_files[label].extend(train_files)\n",
        "  all_eval_example_files[label].extend(eval_files)\n",
        "  if progress is not None:\n",
        "    word_count += 1\n",
        "    progress_bar.update(progress(100 * word_count/len(all_word_list)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "c9Ry4fjoBfsD"
      },
      "source": [
        "In the following cells the wav files from both evaluation and training sets are:\n",
        " * Opened and decoded.\n",
        " * Loudness normalized.\n",
        " * Passed through the embedding model to create embeddings.\n",
        " * Added to a data structure that let's us change the balance between negative, silence (if present) and labeled outputs.\n",
        "\n",
        "\n",
        " resulting in two objects: *eval_data* and *train_data*.\n",
        "\n",
        "\n",
        " The two parameters to consider here are:\n",
        "  * **examples_per_word**: The number examples for each target word that should be loaded. A higher number for the training data will lead to a better model, but it will also take longer to load/train. A good starting point is 40. Small numbers for the eval data may result in easy / hard eval subsets that could give an incorrect impression of the model quality.\n",
        "  * **negatives_multiplier**: How many more non target examples should be loaded. This is set to 25 by default as the speech commands dataset maps 25 words to negative. Also applies to silence examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "A2y7Bmj3Wt7j"
      },
      "outputs": [],
      "source": [
        "#@title Load evaluation set wav.\n",
        "#@markdown Set examples_per_word to \u003e 500 and negatives_multiplier to 25 \n",
        "#@markdown to ensure you load the whole eval set.\n",
        "examples_per_word =  50#@param {type:\"integer\"}\n",
        "negatives_multiplier =  25#@param {type:\"integer\"}\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "\n",
        "print(\"loading eval data\")\n",
        "eval_data = EmbeddingDataFileList(\n",
        "    all_eval_example_files, data_dest_dir, label_max=examples_per_word,\n",
        "    negative_multiplier=negatives_multiplier,\n",
        "    targets=MODEL_LABELS, embedding_model=speech_embedding_model,\n",
        "    progress_bar=progress_bar)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "QO0IJqsh07i4"
      },
      "outputs": [],
      "source": [
        "#@title Load random speech commands wav files for training.\n",
        "#@markdown Set examples_per_word to \u003e 4000 and negatives_multiplier to 25 \n",
        "#@markdown to ensure you load the whole training set.\n",
        "examples_per_word =  50#@param {type:\"integer\"}\n",
        "negatives_multiplier =  25#@param {type:\"integer\"}\n",
        "\n",
        "progress_bar = IPython.display.display(progress(0, 100), display_id=True)\n",
        "\n",
        "print(\"loading train data\")\n",
        "train_data = EmbeddingDataFileList(\n",
        "    all_train_example_files, data_dest_dir, label_max=examples_per_word,\n",
        "    negative_multiplier=negatives_multiplier,\n",
        "    targets=MODEL_LABELS, embedding_model=speech_embedding_model,\n",
        "    progress_bar=progress_bar)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ipFhC0H14Q9v"
      },
      "source": [
        "## Train and Evaluate a Head Model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "iUhwy_1bCg3P"
      },
      "outputs": [],
      "source": [
        "#@title Rebalance and filter data.\n",
        "\n",
        "#@markdown **Labeled_weight** and **negatives_weight** are used to control the ratio of labeled data\n",
        "#@markdown and negative data shown to the model during training and evaluation.\n",
        "\n",
        "#@markdown For every *labeled_weight* keyword examples the model is trained on, \n",
        "#@markdown it is also  trained on *negatives_weight* non keyword examples. \n",
        "#@markdown During rebalancing examples are duplicated to ensure that this ratio holds.\n",
        "\n",
        "labeled_weight = 8 #@param {type:\"slider\", min:1, max:25, step:1}\n",
        "negatives_weight = 1 #@param {type:\"slider\", min:1, max:25, step:1}\n",
        "\n",
        "#@markdown We assume that the keyphrase is spoken in roughly the middle\n",
        "#@markdown of the loaded audio clips. With **context_size** we can choose the \n",
        "#@markdown number of embeddings around the middle to use as a model input.\n",
        "context_size = 16 #@param {type:\"slider\", min:1, max:28, step:1}\n",
        "\n",
        "\n",
        "filter_fn = functools.partial(cut_middle_frame, num_frames=context_size, flatten=False)\n",
        "\n",
        "\n",
        "eval_data.full_rebalance(negatives=negatives_weight, labeled=labeled_weight)\n",
        "all_eval_data = eval_data.get_all_data_shuffled(filter_fn=filter_fn)\n",
        "\n",
        "train_data.full_rebalance(negatives=negatives_weight, labeled=labeled_weight)\n",
        "all_train_data = train_data.get_all_data_shuffled(filter_fn=filter_fn)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "bqtJk_XN9Tcb"
      },
      "outputs": [],
      "source": [
        "#@title Run training and evaluation\n",
        "head_model = \"Convolutional\" #@param [\"Convolutional\", \"Fully_Connected\"] {type:\"string\"}\n",
        "\n",
        "#@markdown Suggested **learning_rate** range 0.00001 - 0.01.\n",
        "learning_rate = 0.001 #@param {type:\"number\"}\n",
        "batch_size = 32\n",
        "#@markdown **epochs_per_eval** and **train_eval_loops** control how long the\n",
        "#@markdown the model is trained. An epoch is defined as the model having seen\n",
        "#@markdown each example at least once, with some examples twice to ensure the\n",
        "#@markdown correct labeled / negatives balance.\n",
        "\n",
        "epochs_per_eval = 1 #@param {type:\"slider\", min:1, max:15, step:1}\n",
        "train_eval_loops = 15 #@param {type:\"slider\", min:5, max:80, step:5}\n",
        "\n",
        "if head_model == \"Convolutional\":\n",
        "  model_fn = functools.partial(_conv_head_model_fn, context=context_size)\n",
        "else:\n",
        "  model_fn = _fully_connected_model_fn\n",
        "\n",
        "\n",
        "trainer = HeadTrainer(model_fn=model_fn,\n",
        "                      input_shape=[context_size,1,96],\n",
        "                      num_targets=len(MODEL_LABELS),\n",
        "                      head_learning_rate=learning_rate,\n",
        "                      batch_size=batch_size)\n",
        "\n",
        "data_trained_on = 0\n",
        "data = [] \n",
        "train_results = []\n",
        "eval_results = []\n",
        "max_data = len(all_train_data) * epochs_per_eval * train_eval_loops + 10\n",
        "\n",
        "\n",
        "def plot_step(plot, max_data, data, train_results, eval_results):\n",
        "  plot.clf()\n",
        "  plot.xlim(0, max_data)\n",
        "  plot.ylim(0.85, 1.05)\n",
        "  plot.plot(data, train_results, \"bo\")\n",
        "  plot.plot(data, train_results, \"b\", label=\"train_results\")\n",
        "  if eval_results:\n",
        "    plot.plot(data, eval_results, \"ro\")\n",
        "    plot.plot(data, eval_results, \"r\", label=\"eval_results\")\n",
        "  plot.legend(loc='lower right', fontsize=24)\n",
        "  plot.xlabel('number of examples trained on', fontsize=22)\n",
        "  plot.ylabel('Accuracy', fontsize=22)\n",
        "  plot.xticks(fontsize=20)\n",
        "  plot.yticks(fontsize=20) \n",
        "\n",
        "plt.figure(figsize=(25, 7))\n",
        "for loop in range(train_eval_loops):\n",
        "  train_accuracy, loss = trainer.epoch_train(all_train_data,\n",
        "                                             epochs=epochs_per_eval)\n",
        "  train_results.append(train_accuracy)\n",
        "  if all_eval_data:\n",
        "    eval_accuracy, loss = trainer.test(all_eval_data)\n",
        "    eval_results.append(eval_accuracy)\n",
        "  else:\n",
        "    eval_results = None\n",
        "\n",
        "  data_trained_on += len(all_train_data) * epochs_per_eval\n",
        "  data.append(data_trained_on)\n",
        "  plot_step(plt, max_data, data, train_results, eval_results)\n",
        "\n",
        "  IPython.display.display(plt.gcf())\n",
        "  if all_eval_data:\n",
        "    print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n",
        "  IPython.display.clear_output(wait=True)\n",
        "\n",
        "if all_eval_data:\n",
        "  print(\"Highest eval accuracy: %.2f percent.\" % (100 * max(eval_results)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fxa9wsJKeiv9"
      },
      "source": [
        "# Export and reuse the head model\n",
        "The following cells show how the head model can be exported and reused in a graph "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "s2GxL706F-BD"
      },
      "outputs": [],
      "source": [
        "#@title Save the head model\n",
        "\n",
        "head_model_module_dir = \"head_model_module_fc_all_data\"\n",
        "trainer.save_head_model(head_model_module_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "r6PfTOh7HIt1"
      },
      "outputs": [],
      "source": [
        "#@title FullModelWrapper - Example Class\n",
        "\n",
        "class FullModelWrapper(object):\n",
        "  \"\"\"A loads a save model classifier.\"\"\"\n",
        "  def __init__(self, embedding_model_dir, head_model_dir):\n",
        "    self._graph = tf.Graph()\n",
        "    self._sess = tf.Session(graph=self._graph)\n",
        "    with self._graph.as_default():\n",
        "      self._samples = tf.placeholder(\n",
        "          tf.float32, shape=[1, None], name='audio_samples')\n",
        "      module_spec = hub.create_module_spec(\n",
        "          module_fn=self._get_module_fn(embedding_model_dir, head_model_dir))\n",
        "      self._module = hub.Module(module_spec, trainable=True)\n",
        "      self._predictions = self._module(self._samples)\n",
        "      with self._sess.as_default():\n",
        "        self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        " \n",
        "  def _get_module_fn(self, embedding_model_dir, head_model_module_dir):\n",
        "    \"\"\"Wraps the model_fn in a tf hub module.\"\"\"\n",
        "    def module_fn():\n",
        "      samples =  tf.placeholder(\n",
        "          tf.float32, shape=[1, None], name='audio_samples')\n",
        "      embedding_module_spec = hub.load_module_spec(embedding_model_dir)\n",
        "      embedding_module = hub.Module(embedding_module_spec)\n",
        "      head_module_spec = hub.load_module_spec(head_model_module_dir)\n",
        "      emb = embedding_module(samples)\n",
        "      head_module = hub.Module(head_module_spec)\n",
        "      logits = head_module(emb)\n",
        "      predictions = tf.nn.softmax(logits)\n",
        "      hub.add_signature(name='default', inputs=samples, outputs=predictions)\n",
        "    return module_fn\n",
        "\n",
        "  def save_head_model(self, save_directory):\n",
        "    \"\"\"Saves the model.\"\"\"\n",
        "    with self._graph.as_default():\n",
        "      self._module.export(save_directory, self._sess)\n",
        "\n",
        "  def infer(self, samples):\n",
        "    samples = samples.reshape((1, -1))\n",
        "    output = self._sess.run(\n",
        "        [self._predictions],\n",
        "        feed_dict={self._samples: samples})\n",
        "    return output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "gS9gCV8SKIfe"
      },
      "outputs": [],
      "source": [
        "#@title Test the full model on zeros\n",
        "full_model = FullModelWrapper(embedding_model_url, head_model_module_dir)\n",
        "full_model.infer(np.zeros((1,32000)))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "JndnmDMp66FL"
      ],
      "name": "With silence - Speech Embeddings 1: Train on the speech commands data.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
